# backtest.py
import backtrader as bt
import pandas as pd
import pickle
from datetime import datetime


class MomentumStrategy(bt.Strategy):
    params = (
        ('top_n', 10),  # 每期持有股票数量
        ('rebalance_days', 20),  # 调仓周期
    )

    def __init__(self):
        self.day_count = 0
        # 加载预先生成的交易信号
        with open('trading_signals.pkl', 'rb') as f:
            self.trading_signals = pickle.load(f)

        # 打印信号信息用于调试
        print(f"加载的交易信号日期范围: {list(self.trading_signals.keys())[:5]} ...")
        print(f"信号总数: {len(self.trading_signals)}")

    def next(self):
        self.day_count += 1

        # 到达调仓日才执行调仓
        if self.day_count % self.params.rebalance_days != 0:
            return

        current_date = self.datas[0].datetime.date(0)
        current_date_str = current_date.strftime('%Y-%m-%d')

        print(f"调仓日: {current_date_str}, 检查信号...")

        if current_date_str not in self.trading_signals:
            print(f"日期 {current_date_str} 没有交易信号")
            return

        # 获取当前调仓日的信号
        target_stocks = self.trading_signals[current_date_str]
        print(f"目标股票: {target_stocks}")

        # 清空当前持仓
        for data in self.datas:
            position = self.getposition(data)
            if position.size > 0:
                print(f"平仓: {data._name}, 数量: {position.size}")
                self.close(data)

        # 等权重买入新的目标股票
        target_datas = [d for d in self.datas if d._name in target_stocks]
        if not target_datas:
            print("没有找到对应的股票数据")
            return

        weight = 1.0 / len(target_datas)
        print(f"买入 {len(target_datas)} 只股票, 每只权重: {weight:.2%}")

        for data in target_datas:
            self.order_target_percent(data, target=weight)
            print(f"下单买入: {data._name}, 权重: {weight:.2%}")


def run_backtest():
    # 创建Cerebro引擎
    cerebro = bt.Cerebro()

    # 加载数据
    try:
        with open('stock_data.pkl', 'rb') as f:
            all_data = pickle.load(f)
        print(f"加载了 {len(all_data)} 只股票的数据")
    except FileNotFoundError:
        print("错误: 找不到 stock_data.pkl 文件")
        return [], cerebro
    except Exception as e:
        print(f"加载数据错误: {e}")
        return [], cerebro

    # 检查交易信号文件
    try:
        with open('trading_signals.pkl', 'rb') as f:
            trading_signals = pickle.load(f)
        print(f"交易信号日期: {list(trading_signals.keys())[:5]} ...")
    except FileNotFoundError:
        print("错误: 找不到 trading_signals.pkl 文件")
        return [], cerebro

    # 为每只股票创建数据源
    added_stocks = 0
    for ts_code, df in all_data.items():
        try:
            # 确保数据格式符合backtrader要求
            data_df = df[['open', 'high', 'low', 'close', 'vol']].copy()
            data_df.columns = ['open', 'high', 'low', 'close', 'volume']
            data_df.index = pd.to_datetime(data_df.index)

            # 检查数据是否足够
            if len(data_df) < 100:
                continue

            # 创建数据feed
            data_feed = bt.feeds.PandasData(
                dataname=data_df,
                name=ts_code  # 为数据命名，便于在策略中识别
            )
            cerebro.adddata(data_feed)
            added_stocks += 1

        except Exception as e:
            print(f"添加股票 {ts_code} 数据时出错: {e}")
            continue

    print(f"成功添加 {added_stocks} 只股票数据到回测引擎")

    if added_stocks == 0:
        print("错误: 没有添加任何股票数据")
        return [], cerebro

    # 添加策略
    cerebro.addstrategy(MomentumStrategy)

    # 设置初始资金
    cerebro.broker.setcash(100000.0)

    # 设置佣金
    cerebro.broker.setcommission(commission=0.001)  # 0.1%佣金

    # 添加分析器
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')
    cerebro.addanalyzer(bt.analyzers.Returns, _name='returns')
    cerebro.addanalyzer(bt.analyzers.AnnualReturn, _name='annual')
    cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='trades')
    cerebro.addanalyzer(bt.analyzers.Transactions, _name='transactions')

    # 运行回测
    print('初始投资组合价值: %.2f' % cerebro.broker.getvalue())

    try:
        results = cerebro.run()
        print('最终投资组合价值: %.2f' % cerebro.broker.getvalue())

        if not results:
            print("回测没有返回任何结果")
            return [], cerebro

        return results, cerebro

    except Exception as e:
        print(f"回测运行错误: {e}")
        return [], cerebro


if __name__ == "__main__":
    results, cerebro = run_backtest()

    if not results:
        print("回测失败，无法继续分析")
        exit(1)

    # 修复这里：results 是列表，需要获取第一个元素
    strategy_instance = results[0]

    # 检查是否有交易发生
    try:
        trade_analysis = strategy_instance.analyzers.trades.get_analysis()
        if not trade_analysis:
            print("回测期间没有执行任何交易")
            print("可能的原因:")
            print("1. 交易信号日期与数据日期不匹配")
            print("2. 调仓周期设置过长")
            print("3. 目标股票在数据中不存在")
        else:
            print(f"交易统计: 总交易次数 {trade_analysis.get('total', {}).get('total', 0)}")
    except Exception as e:
        print(f"交易分析错误: {e}")

    # 打印分析结果
    print("\n=== 回测结果分析 ===")

    try:
        sharpe_analysis = strategy_instance.analyzers.sharpe.get_analysis()
        sharpe_ratio = sharpe_analysis.get('sharperatio', '无法计算')
        print("夏普比率:", sharpe_ratio)
    except Exception as e:
        print(f"夏普比率计算错误: {e}")

    try:
        drawdown_analysis = strategy_instance.analyzers.drawdown.get_analysis()
        max_drawdown = drawdown_analysis.get('max', {}).get('drawdown', '无法计算')
        if isinstance(max_drawdown, float):
            print("最大回撤:", f"{max_drawdown:.2%}")
        else:
            print("最大回撤:", max_drawdown)
    except Exception as e:
        print(f"最大回撤计算错误: {e}")

    try:
        returns_analysis = strategy_instance.analyzers.returns.get_analysis()
        total_return = returns_analysis.get('rtot', '无法计算')
        if isinstance(total_return, float):
            print("总收益率:", f"{total_return:.2%}")
        else:
            print("总收益率:", total_return)
    except Exception as e:
        print(f"总收益率计算错误: {e}")

    try:
        annual_analysis = strategy_instance.analyzers.annual.get_analysis()
        print("年化收益率:", annual_analysis)
    except Exception as e:
        print(f"年化收益率计算错误: {e}")

    # 绘制回测结果
    try:
        cerebro.plot(style='candlestick')
    except Exception as e:
        print(f"绘图错误: {e}")